from contextlib import contextmanager
from math import sqrt

from khronos.utils import Namespace

# ---------------------------------------------------------
# Data collection functions -------------------------------
def mk_collect_fnc(name, aggregator):
    @contextmanager
    def collect_fnc(indicator_fnc, collector_fnc):
        x0 = indicator_fnc()
        yield
        x1 = indicator_fnc()
        x = aggregator(x0, x1)
        collector_fnc(x)
    collect_fnc.__name__ = "collect_" + name
    return collect_fnc
    
@contextmanager
def collect_both(indicator_fnc, collector_fnc):
    collector_fnc(indicator_fnc())
    yield
    collector_fnc(indicator_fnc())
    
# A collection of context managers to automatically collect data into stat collectors
collect = Namespace(diff=mk_collect_fnc("diff", lambda x0, x1: (x1 - x0)),
                    sum=mk_collect_fnc("sum", lambda x0, x1: (x1 + x0)),
                    mean=mk_collect_fnc("mean", lambda x0, x1: (x0 + x1) / 2.0),
                    min=mk_collect_fnc("min", lambda x0, x1: min(x0, x1)),
                    max=mk_collect_fnc("max", lambda x0, x1: max(x0, x1)),
                    both=collect_both)
                    
# ---------------------------------------------------------
# Sampling functions --------------------------------------
def mk_sample_fnc(condition, name=None):
    def sample_fnc(distr, *args, **kwargs):
        try:
            tries_left = tries = kwargs["tries"]
        except KeyError:
            tries_left = tries = 100
        else:
            del kwargs["tries"]
        while True:
            if tries_left <= 0:
                raise ValueError("unable to get valid sample in %d tries" % (tries,))
            value = distr(*args, **kwargs)
            if condition(value):
                return value
            tries_left -= 1
    if name is None:
        name = condition.__name__
    sample_fnc.__name__ = "sample_" + name
    return sample_fnc
    
# A collection of functions to sample probability distributions until a    
# sample that meets a certain condition (e.g. nonnegative) is observed
sample = Namespace(positive=mk_sample_fnc(lambda x: x > 0.0, "pos"), 
                   negative=mk_sample_fnc(lambda x: x < 0.0, "neg"), 
                   nonpositive=mk_sample_fnc(lambda x: x <= 0.0, "nonpos"), 
                   nonnegative=mk_sample_fnc(lambda x: x >= 0.0, "nonneg"),
                   conditional=mk_sample_fnc)

# ---------------------------------------------------------
# Basic descriptive statistics ----------------------------
def mean(values):
    return sum(values) / float(len(values))
    
def variance(values, mu=None, bessel_correction=True):
    if mu is None:
        mu = mean(values)
    n = len(values) - (1 if bessel_correction else 0)
    return float(sum((mu - x)**2 for x in values)) / n
    
def stddev(values, mu=None, bessel_correction=True):
    return sqrt(variance(values, mu, bessel_correction))
    
def quantile(q, values):
    """Quantile calculation for a given iterable with real numbers. The quantile argument 'q' 
    must belong to the interval ]0, 1].
    NOTE: the behavior of this function is not satisfactory yet. I need to read more about 
    quantile calculation and perhaps implement a different method."""
    if not 0.0 < q <= 1.0:
        raise ValueError("quantile must lie in ]0, 1]")
    values = sorted(values)
    index = q * (len(values) + 1) - 1
    if index < 0.0:
        raise ValueError("quantile undefined for %f" % (q,))
    idx0 = int(index)
    idx1 = idx0 + 1
    diff = index - int(index)
    if diff == 0.0:
        return values[idx0]
    return (1.0 - diff) * values[idx0] + diff * values[idx1]
    
def median(values):
    return quantile(0.5, values)
    